import fs from 'fs/promises';
import { Translator } from './dist/services/translator.js';
import pLimit from 'p-limit';

// Function to recursively extract all string values from the JSON object
function extractStrings(obj, strings) {
  if (typeof obj === 'string') {
    // Filter out non-Chinese strings and very short strings
    if (/[\u4e00-\u9fa5]/.test(obj) && obj.trim().length > 2) {
      strings.add(obj.trim());
    }
  } else if (Array.isArray(obj)) {
    obj.forEach(item => extractStrings(item, strings));
  } else if (typeof obj === 'object' && obj !== null) {
    Object.values(obj).forEach(value => extractStrings(value, strings));
  }
}

async function createTranslationMemory() {
  try {
    console.log("Reading the source JSON file...");
    const jsonContent = await fs.readFile('gsi_drp4dvar_wiki_outline.json', 'utf-8');
    const jsonData = JSON.parse(jsonContent);

    console.log("Extracting Chinese strings for translation...");
    const stringsToTranslate = new Set();
    extractStrings(jsonData, stringsToTranslate);

    const translator = new Translator();
    const limit = pLimit(2); // Set concurrency to 2 as requested
    const batchSize = 20; // Group 20 sentences per API call

    const sourceTexts = Array.from(stringsToTranslate);
    const batches = [];
    for (let i = 0; i < sourceTexts.length; i += batchSize) {
      batches.push(sourceTexts.slice(i, i + batchSize));
    }

    console.log(`Found ${sourceTexts.length} unique strings, grouped into ${batches.length} batches of up to ${batchSize} sentences each.`);

    const translationPromises = batches.map((batch, batchIndex) =>
      limit(async () => {
        const batchText = batch.join('\n'); // Use a clear separator
        console.log(`Translating batch ${batchIndex + 1}/${batches.length}...`);
        try {
          const translatedBatchText = await translator.translate(batchText, 'English');
          const translatedSentences = translatedBatchText.split('\n');

          if (batch.length === translatedSentences.length) {
            return batch.map((source, i) => ({
              source: source,
              target: translatedSentences[i].trim()
            }));
          } else {
            console.error(` - Mismatch in batch ${batchIndex + 1}: source has ${batch.length}, translation has ${translatedSentences.length}. Skipping batch.`);
            return [];
          }
        } catch (error) {
          console.error(` - Error translating batch ${batchIndex + 1}:`, error);
          return [];
        }
      })
    );

    const results = await Promise.all(translationPromises);
    
    // Flatten the array of arrays and filter out any invalid entries
    const translationMemory = results.flat().filter(item => item && item.target);

    console.log("Writing the new translation memory file...");
    await fs.writeFile('translation-memory.json', JSON.stringify(translationMemory, null, 2), 'utf-8');

    console.log(`Successfully created the detailed translation-memory.json file with ${translationMemory.length} entries.`);

  } catch (error) {
    console.error("An error occurred during the process:", error);
  }
}

createTranslationMemory();

